package com.srtp.secureprotocol.cipherstream;

import com.srtp.secureprotocol.util.Util;
import org.whispersystems.libsignal.InvalidMacException;
import org.whispersystems.libsignal.InvalidMessageException;

import java.io.File;
import java.io.FileInputStream;
import java.io.FilterInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.Arrays;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.Mac;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.ShortBufferException;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

/**
 * 加密文件输入流，继承自FilterInputStream
 * 提供文件解密功能，支持AES/CBC/PKCS5Padding加密模式
 * 包含MAC（消息认证码）验证确保数据完整性
 */
public class CipherInputStream extends FilterInputStream {

    // 加密相关常量
    private static final int BLOCK_SIZE = 16;      // AES块大小（字节）
    private static final int CIPHER_KEY_SIZE = 32;      // 加密密钥长度（字节）
    private static final int MAC_KEY_SIZE = 32;      // MAC密钥长度（字节）

    // 解密状态变量
    private Cipher cipher;          // AES解密器
    private boolean done;            // 解密完成标志
    private long totalDataSize;   // 待解密数据总大小
    private long totalRead;       // 已读取数据量
    private byte[] overflowBuffer;  // 解密溢出缓冲区

    /**
     * 为加密文件创建解密输入流
     *
     * @param file                加密文件
     * @param plaintextLength     明文长度（0表示未知）
     * @param combinedKeyMaterial 组合密钥材料（加密密钥+MAC密钥）
     * @param digest              预期的摘要值（用于完整性校验）
     * @return 配置好的解密输入流
     * @throws InvalidMessageException 消息格式无效
     * @throws IOException             文件读写错误
     */
    public static InputStream createForFile(File file, long plaintextLength, byte[] combinedKeyMaterial, byte[] digest)
            throws InvalidMessageException, IOException {
        try {
            // 分割组合密钥为加密密钥和MAC密钥
            byte[][] parts = Util.split(combinedKeyMaterial, CIPHER_KEY_SIZE, MAC_KEY_SIZE);

            // 初始化HMAC-SHA256验证器
            Mac mac = Mac.getInstance("HmacSHA256");
            mac.init(new SecretKeySpec(parts[1], "HmacSHA256"));

            // 验证文件长度有效性
            if (file.length() <= BLOCK_SIZE + mac.getMacLength()) {
                throw new InvalidMessageException("文件长度小于加密开销!");
            }

            // 检查摘要参数
            if (digest == null) {
                throw new InvalidMacException("缺少摘要参数!");
            }

            // 验证MAC完整性
            try (FileInputStream fin = new FileInputStream(file)) {
                verifyMac(fin, file.length(), mac, digest);
            }

            // 创建解密流（排除IV和MAC部分）
            InputStream inputStream = new CipherInputStream(
                    new FileInputStream(file),
                    parts[0],
                    file.length() - BLOCK_SIZE - mac.getMacLength()
            );

            // 如果指定了明文长度，添加长度限制流
            if (plaintextLength != 0) {
                inputStream = new ContentLengthInputStream(inputStream, plaintextLength);
            }

            return inputStream;
        } catch (NoSuchAlgorithmException | InvalidKeyException e) {
            throw new AssertionError(e); // 不应发生的加密算法错误
        } catch (InvalidMacException e) {
            throw new InvalidMessageException(e); // MAC验证失败
        }
    }

    /**
     * 解密流构造函数（私有）
     *
     * @param inputStream   底层输入流（包含加密数据）
     * @param cipherKey     加密密钥
     * @param totalDataSize 待解密数据总大小
     * @throws IOException 流读取错误
     */
    private CipherInputStream(InputStream inputStream, byte[] cipherKey, long totalDataSize)
            throws IOException {
        super(inputStream);

        try {
            // 读取初始化向量(IV)
            byte[] iv = new byte[BLOCK_SIZE];
            readFully(iv);

            // 初始化AES/CBC/PKCS5Padding解密器
            this.cipher = Cipher.getInstance("AES/CBC/PKCS5Padding");
            this.cipher.init(Cipher.DECRYPT_MODE,
                    new SecretKeySpec(cipherKey, "AES"),
                    new IvParameterSpec(iv)
            );

            // 初始化状态变量
            this.done = false;
            this.totalRead = 0;
            this.totalDataSize = totalDataSize;
        } catch (NoSuchAlgorithmException | InvalidKeyException |
                 NoSuchPaddingException | InvalidAlgorithmParameterException e) {
            throw new AssertionError(e); // 加密配置错误
        }
    }

    @Override
    public int read(byte[] buffer) throws IOException {
        return read(buffer, 0, buffer.length);
    }

    @Override
    public int read(byte[] buffer, int offset, int length) throws IOException {
        // 根据解密状态选择读取策略
        if (totalRead < totalDataSize) {
            return readIncremental(buffer, offset, length); // 增量解密
        } else if (!done) {
            return readFinal(buffer, offset, length);      // 最终解密块
        } else {
            return -1; // 数据已全部读取
        }
    }

    @Override
    public boolean markSupported() {
        return false; // 不支持标记/重置
    }

    @Override
    public long skip(long byteCount) throws IOException {
        // 通过实际读取实现跳过（因解密流特性）
        long skipped = 0L;
        while (skipped < byteCount) {
            byte[] buf = new byte[Math.min(4096, (int) (byteCount - skipped))];
            int read = read(buf);
            if (read == -1) break;
            skipped += read;
        }
        return skipped;
    }

    /**
     * 处理最终解密块（PKCS5填充移除）
     */
    private int readFinal(byte[] buffer, int offset, int length) throws IOException {
        try {
            // 执行最终解密操作（处理填充）
            int flourish = cipher.doFinal(buffer, offset);
            done = true; // 标记解密完成
            return flourish;
        } catch (IllegalBlockSizeException | BadPaddingException | ShortBufferException e) {
            throw new IOException("最终解密失败", e);
        }
    }

    /**
     * 增量解密读取（处理大部分数据）
     */
    private int readIncremental(byte[] buffer, int offset, int length) throws IOException {
        int readLength = 0;

        // 处理上次溢出的解密数据
        if (overflowBuffer != null) {
            if (overflowBuffer.length > length) {
                System.arraycopy(overflowBuffer, 0, buffer, offset, length);
                overflowBuffer = Arrays.copyOfRange(overflowBuffer, length, overflowBuffer.length);
                return length;
            } else if (overflowBuffer.length == length) {
                System.arraycopy(overflowBuffer, 0, buffer, offset, length);
                overflowBuffer = null;
                return length;
            } else {
                System.arraycopy(overflowBuffer, 0, buffer, offset, overflowBuffer.length);
                readLength = overflowBuffer.length;
                offset += readLength;
                length -= readLength;
                overflowBuffer = null;
            }
        }

        // 调整读取长度避免超出数据范围
        if (length + totalRead > totalDataSize) {
            length = (int) (totalDataSize - totalRead);
        }

        // 从底层流读取加密数据（预留块空间）
        byte[] internalBuffer = new byte[length];
        int maxRead = internalBuffer.length <= cipher.getBlockSize() ?
                internalBuffer.length :
                internalBuffer.length - cipher.getBlockSize();
        int read = super.read(internalBuffer, 0, maxRead);

        // 文件结束检查
        if (read == -1) {
            if (totalRead < totalDataSize) throw new IOException("文件过早结束");
            return -1;
        }

        totalRead += read;

        try {
            // 预估解密后数据大小
            int outputLen = cipher.getOutputSize(read);

            // 情况1：解密数据可直接放入目标缓冲区
            if (outputLen <= length) {
                int decrypted = cipher.update(internalBuffer, 0, read, buffer, offset);
                readLength += decrypted;
                return readLength;
            }

            // 情况2：解密数据超出目标缓冲区容量
            byte[] transientBuffer = new byte[outputLen];
            outputLen = cipher.update(internalBuffer, 0, read, transientBuffer, 0);

            if (outputLen <= length) {
                System.arraycopy(transientBuffer, 0, buffer, offset, outputLen);
                readLength += outputLen;
            } else {
                // 部分数据溢出到下次读取
                System.arraycopy(transientBuffer, 0, buffer, offset, length);
                overflowBuffer = Arrays.copyOfRange(transientBuffer, length, outputLen);
                readLength += length;
            }
            return readLength;
        } catch (ShortBufferException e) {
            throw new AssertionError("缓冲区不足", e);
        }
    }

    /**
     * 验证消息认证码（MAC）和文件摘要
     *
     * @param inputStream 输入流（包含加密数据+MAC）
     * @param length      文件总长度
     * @param mac         初始化的MAC验证器
     * @param theirDigest 预期的摘要值
     * @throws InvalidMacException MAC或摘要验证失败
     */
    private static void verifyMac(InputStream inputStream, long length, Mac mac, byte[] theirDigest)
            throws InvalidMacException {
        try {
            // 初始化SHA-256摘要器
            MessageDigest digest = MessageDigest.getInstance("SHA256");
            int remainingData = Util.toIntExact(length) - mac.getMacLength(); // 排除MAC尾部的数据长度
            byte[] buffer = new byte[4096];

            // 读取并处理整个数据部分
            while (remainingData > 0) {
                int read = inputStream.read(buffer, 0, Math.min(buffer.length, remainingData));
                if (read == -1) throw new IOException("文件过早结束");

                // 更新MAC和摘要
                mac.update(buffer, 0, read);
                digest.update(buffer, 0, read);
                remainingData -= read;
            }

            // 计算并验证MAC
            byte[] ourMac = mac.doFinal();
            byte[] theirMac = new byte[mac.getMacLength()];
            Util.readFully(inputStream, theirMac); // 读取文件中的MAC值

            if (!MessageDigest.isEqual(ourMac, theirMac)) {
                throw new InvalidMacException("MAC值不匹配!");
            }

            // 计算并验证摘要
            byte[] ourDigest = digest.digest(theirMac); // 对MAC值做摘要

            if (theirDigest != null && !MessageDigest.isEqual(ourDigest, theirDigest)) {
                throw new InvalidMacException("摘要值不匹配!");
            }

        } catch (IOException | ArithmeticException e1) {
            throw new InvalidMacException(e1);
        } catch (NoSuchAlgorithmException e) {
            throw new AssertionError(e);
        }
    }

    /**
     * 完全填充缓冲区（确保读取指定字节数）
     *
     * @param buffer 目标缓冲区
     * @throws IOException 读取不足或流结束
     */
    private void readFully(byte[] buffer) throws IOException {
        int offset = 0;
        while (offset < buffer.length) {
            int read = super.read(buffer, offset, buffer.length - offset);
            if (read == -1) throw new IOException("流过早结束");
            offset += read;
        }
    }
}